from core import *
from OpenGL.GL.framebufferobjects import *
from types import Vector3, Vector4
from numpy import float32
import Image

MAX_TEXTURES = 8

ATTACHMENTS = range(GL_COLOR_ATTACHMENT0, GL_COLOR_ATTACHMENT8 + 1)

''' We will use a texture "stack", each texture unit acts as a stack, binding adds the texture
on top of the stack, unbinding it pops it off, binding the one below it '''
Stacks = [[] for x in range(MAX_TEXTURES)]

class Texture(GlObject):
    def __init__(self, width = 256, height = 256, format = GL_RGBA8, mode = GL_RGBA, type = GL_UNSIGNED_BYTE, mag_filter = GL_LINEAR, min_filter = GL_NEAREST, wrap = GL_REPEAT, data = None, mipmap = False):
        GlObject.__init__(self, glGenTextures(1))
        glBindTexture(GL_TEXTURE_2D, self.handle)

        if mipmap:
            glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR)
            glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, mag_filter)
            gluBuild2DMipmaps( GL_TEXTURE_2D, format, width, height, mode, type, data )
        else:
            glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MIN_FILTER, min_filter)
            glTexParameteri(GL_TEXTURE_2D,GL_TEXTURE_MAG_FILTER, mag_filter)
            glTexImage2D(GL_TEXTURE_2D, 0, format, width, height, 0, mode, type, data)
        
        glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, wrap)
        glTexParameteri( GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, wrap)
        
        self.slot = None
        
    @staticmethod
    def FromImage(path, **kwargs):
        '''
            Creates a texture from image file
        '''
        im = Image.open(path)
        
        try:
            # get image meta-data (dimensions) and data
            ix, iy, image, channels = im.size[0], im.size[1], im.tostring("raw", "RGBA", 0, -1), 4
        except SystemError:
            # has no alpha channel, synthesize one, see the
            # texture module for more realistic handling
            ix, iy, image, channels = im.size[0], im.size[1], im.tostring("raw", "RGBX", 0, -1), 4

        # Create Texture
        return Texture(ix, iy, format = GL_RGBA, data = image, **kwargs)
        
    def bind(self, slot):
        glActiveTexture(GL_TEXTURE0 + slot)
        self.slot = slot
        Stacks[slot].append(self)
        glBindTexture(GL_TEXTURE_2D, self.handle)
        if len(Stacks[slot]) > 8:
            raise Exception('Texture Stack Overflow!')
        
    def unbind(self):
        ''' 
            Note this function may behave slightly oddly.
            By calling it, there is no guarantee the texture being unbound is this one.
        '''
        if self.slot is not None:
            glActiveTexture(GL_TEXTURE0 + self.slot)
            stack = Stacks[self.slot]
            stack.pop()
            
            # bind the next one
            if len(stack) == 0:
                glBindTexture(GL_TEXTURE_2D, 0)
            else:
                glBindTexture(GL_TEXTURE_2D, stack[-1].handle)
        self.slot = None
        
    def __repr__(self):
        return 'T'

class NoTexture(Texture):
    '''
        Empty texture, use when a Texture object is required but
        is not available
    '''
    def __init__(self):
        pass
        
    def bind(self, slot):
        pass
        
    def unbind(self):
        pass
        
        
class TextureGroup():
    '''
        Texture groups are array of textures to be bound and unbound at the same time
    '''
    
    def __init__(self, *args):
        self._textures = args

    def bindAll(self, start = 0):
        '''
            Textures bound consecutively, starting at the specified texture unit
        '''
        i = start
        for t in self._textures:
            t.bind(i)
            i += 1
    
    def unbindAll(self):
        for t in self._textures:
            t.unbind()
        
class FBO(GlObject):
    '''
        Frame Buffer Object wrapper - implementation geared towards
        off-screen rendering to texture.
    '''    
    def __init__(self, width = 256, height = 256, depth = True):
        GlObject.__init__(self, glGenFramebuffers(1))
        fbo = self.handle
        glBindFramebuffer(GL_FRAMEBUFFER, fbo)
        
        # adding a depth buffer
        if depth:
            depth = glGenRenderbuffers(1)
            
            glBindRenderbuffer(GL_RENDERBUFFER, depth)
            glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH_COMPONENT24, width, height)
            glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_ATTACHMENT, GL_RENDERBUFFER, depth)
            
        # remember stuff
        self.width = width
        self.height = height
        self.textures = []
        self.targets = 0
        self.textures_enum = None
        
    def addTarget(self, slot = None, **kwargs):
        '''
            Adds a render target (basically a texture). Each FBO must have at least one target
        '''
        img = Texture(self.width, self.height, **kwargs)
        if slot is None:
            slot = self.targets
        glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0 + slot, GL_TEXTURE_2D, img.handle, 0)
        self.textures.append(img)
        
        # update the inner vars for speed
        self.targets = len(self.textures)
        
    def check(self):
        '''
            Check if FBO is ready
        '''
        status = glCheckFramebufferStatus(GL_FRAMEBUFFER)
        assert status == GL_FRAMEBUFFER_COMPLETE, status
        
    def engage(self):
        '''
            Activates rendering to FBO. All subsequent draw calls will 
            render to given target(s). It is up to the developer to 
            render to all target(s)
        '''
        glBindFramebuffer(GL_FRAMEBUFFER, self.handle)
        glPushAttrib(GL_VIEWPORT_BIT)
        glViewport(0,0, self.width, self.height)
        glDrawBuffers(self.targets, ATTACHMENTS[:self.targets])
        
    def disengage(self):
        '''
            Disactivates FBO. Call after engage and all rendering has been finished.
        '''
        glPopAttrib()
        glBindFramebuffer(GL_FRAMEBUFFER, 0)

    def bindTargets(self, *slots):
        '''
            Bind target textures to given slots
        '''
        for i, t in enumerate(self.textures):
            t.bind(slots[i])
            
    def unbindTargets(self):
        for t in self.textures:
            t.unbind()
            
    def render(self):
        '''
            Displays the FBO render result by entering Ortho mode and
            drawing a rect that fills the screen.
        '''
        # clear
        glClearColor(1.0, 0.0, 0.0, 0.0)        # red for ALERT!
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        
        # enter ortho mode
        glMatrixMode(GL_PROJECTION)
        glPushMatrix()
        glLoadIdentity()
        gluOrtho2D(0, self.width, 0, self.height)
        glScalef(1, -1,1)
        glTranslatef(0, -self.height, 0)
        
        glMatrixMode(GL_MODELVIEW)
        
        # Draw a Quad showing the texture
        glLoadIdentity()
        glBegin(GL_QUADS)
        glTexCoord2f(0.0, 1.0)
        glVertex2f(0.0, 0.0)
        glTexCoord2f(0.0, 0.0)
        glVertex2f( 0.0, self.height)
        glTexCoord2f(1.0, 0.0)
        glVertex2f( self.width,  self.height)
        glTexCoord2f(1.0, 1.0)
        glVertex2f(self.width,  0.0)
        glEnd()
        
        # stop ortho view
        glMatrixMode(GL_PROJECTION)
        glPopMatrix()
        glMatrixMode(GL_MODELVIEW)
        
        
class Shader(GlObject):
    VERTEX = GL_VERTEX_SHADER
    FRAGMENT = GL_FRAGMENT_SHADER
    def __init__(self, source, type):
        GlObject.__init__(self, compileShader(source, type))

    @staticmethod
    def fromFile(path, source_mod = lambda x:x):
        ''' 
            also infers the type from ending, use .vert and .frag
            
            Supplying source_mod allows you to modify the source before it gets compiled
        '''
        ext = path[-5:]
        source = source_mod(open(path).read())
        if ext == '.frag':
            return Shader(source, Shader.FRAGMENT)
        elif ext == '.vert':
            return Shader(source, Shader.VERTEX)
        
class ShaderProgram(GlObject):
    uniform_tmap = {
        (int,1): glUniform1i,
        (float,1): glUniform1f,
        (float,3): glUniform3f,
        (float,4): glUniform4f,
        (float32, 4): glUniform4f,
        (float32, 3): glUniform3f,
        (float32, 1): glUniform1f,
    }
    def __init__(self, vert_shader, frag_shader):
        GlObject.__init__(self, compileProgram(
            vert_shader.handle,
            frag_shader.handle,
        ))
    
    def use(self, **uniforms):
        glUseProgram(self.handle)
        self.passUniforms(**uniforms)
    
    def passUniforms(self, **uniforms):
        for k,v in uniforms.items():
            if type(v) is list:
                # array
                try:
                    # array of tuple-type - will fail on v[0][0]
                    tc = self.uniform_tmap[type(v[0][0]), len(v[0])] # get type converter
                    for i in range(len(v)):
                        loc = glGetUniformLocation(self.handle, '%s[%i]' % (k, i))
                        tc(loc, *v[i])
                except TypeError:
                    tc = self.uniform_tmap[type(v[0]), 1]
                    for i in range(len(v)):
                        loc = glGetUniformLocation(self.handle, '%s[%i]' % (k, i))
                        tc(loc, v[i])
            elif type(v) is tuple:
                raise Exception("not yet implemented")
            else:
                loc = glGetUniformLocation(self.handle, k)
                self.uniform_tmap[type(v),1](loc, v)
        ''' passes uniforms (assume all integers) '''
    
    @staticmethod
    def fromFiles(path, source_mod = lambda x:x):
        '''
            Create program from sources path + '.vert' and path + '.frag'
        '''
        return ShaderProgram(Shader.fromFile(path + '.vert', source_mod), Shader.fromFile(path + '.frag', source_mod))
        
    @staticmethod
    def unuse():
        glUseProgram(0)